import time
import threading
import sys
from cushy_serial import CushySerial
from BrainLinkParser import BrainLinkParser
import pygame
import numpy as np

# ============== 配置 ==============
COM_PORT = "COM11"

running = threading.Event()
running.set()

pygame.mixer.init(frequency=44100, size=-16, channels=2, buffer=512)

# ================== 歌曲库（C大调简单旋律） ==================
songs = {
    "1": {"name": "小星星", "notes": [523,523,659,659,698,698,659, 587,587,523,523,493,493,440]},
    "2": {"name": "生日快乐", "notes": [523,523,587,523,698,659, 523,523,587,523,784,698, 523,523,1046,880,698,659,587, 932,932,880,698,784,698]},
    "3": {"name": "欢乐颂", "notes": [659,659,698,740,784,784,698, 659,587,523,523,587,659,659,587,587]},
    "4": {"name": "两只老虎", "notes": [523,587,659,523,523,587,659,523, 659,698,784, 659,698,784, 784,880,784,698,659,523, 523,587,659,523]},
    "5": {"name": "听我说谢谢你", "notes": [523,523,587,659,659,587,523, 659,698,784,880,880,784,698, 659,587,523,523,587,659,523]}   # 简化主旋律
}

current_song_key = "2"   # "生日快乐" 默认小星星
melody_notes = songs[current_song_key]["notes"]
melody_index = 0

current_freq = 440.0
current_volume = 0.5

def play_smooth_note(freq, volume, duration=0.16):
    sample_rate = 44100
    t = np.linspace(0, duration, int(sample_rate * duration), False)
    wave = np.sin(2 * np.pi * freq * t)
    wave += 0.28 * np.sin(2 * np.pi * freq * 2 * t)   # 谐波让声音更丰富
    envelope = np.exp(-2.8 * t)
    audio = (wave * envelope * volume * 22000).astype(np.int16)
    stereo = np.column_stack((audio, audio))
    sound = pygame.sndarray.make_sound(stereo)
    sound.play()

def change_song(song_key):
    global melody_notes, melody_index, current_song_key
    if song_key in songs:
        current_song_key = song_key
        melody_notes = songs[song_key]["notes"]
        melody_index = 0
        print(f"\n🎵 已切换到：{songs[song_key]['name']}   请用意念继续演奏！\n")

# ============== EEG 回调（专注力大幅放大） ==============
def on_eeg(data):
    global melody_index, current_freq, current_volume

    attention = getattr(data, 'attention', 0)
    meditation = getattr(data, 'meditation', 0)
    blink = getattr(data, 'blinkStrength', getattr(data, 'blink', 0))

    print(f"专注度: {attention:3d} | 放松度: {meditation:3d} | 眨眼: {blink:3d}   当前歌曲: {songs[current_song_key]['name']}")
    sys.stdout.flush()

    # 大幅放大专注力影响（即使小波动也有明显变化）
    attention_boost = (attention - 50) * 5.5
    meditation_boost = (meditation - 45) * 2.8

    # 基础频率 + 旋律音符偏移（专注高→明显变亮向上）
    base = melody_notes[melody_index % len(melody_notes)]
    target_freq = base + attention_boost * 1.4 - meditation_boost * 0.9
    target_freq = max(280, min(980, target_freq))

    # 平滑过渡
    current_freq = current_freq * 0.76 + target_freq * 0.24
    current_volume = 0.48 + (attention_boost * 0.009) + (meditation_boost * 0.007)
    current_volume = max(0.38, min(0.98, current_volume))

    play_smooth_note(current_freq, current_volume, duration=0.17)

    melody_index += 1

    # 眨眼加强效果
    if blink > 60:
        play_smooth_note(current_freq + 140, current_volume + 0.28, 0.10)

# 其他回调
def on_extend_eeg(data): pass
def on_gyro(x, y, z): pass
def on_rr(rr1, rr2, rr3): pass
def on_raw(raw): pass

parser = BrainLinkParser(on_eeg, on_extend_eeg, on_gyro, on_rr, on_raw)

# ============== 串口线程 + 歌曲切换菜单 ==============
def start_brainlink():
    serial = None
    try:
        serial = CushySerial(COM_PORT, 115200)
        print(f"✅ 串口 {COM_PORT} 打开成功！")

        @serial.on_message()
        def handle_message(msg: bytes):
            if msg and running.is_set():
                parser.parse(msg)

        print("🎹 Brainlink Pro 脑波音乐已启动！（多首歌曲版）")
        print("可用歌曲：")
        print("   1. 小星星")
        print("   2. 生日快乐")
        print("   3. 欢乐颂")
        print("   4. 两只老虎")
        print("   5. 听我说谢谢你")
        print("\n在程序运行时输入数字（1~5）+ 回车，即可切换歌曲\n")

        # 测试当前歌曲
        print(f"正在播放 {songs[current_song_key]['name']} 测试旋律...")
        for i in range(10):
            play_smooth_note(melody_notes[i % len(melody_notes)], 0.78, 0.26)
            time.sleep(0.30)
        print("测试结束，开始用意念演奏吧！\n")

        while running.is_set():
            # 支持键盘输入切换歌曲（非阻塞）
            if sys.stdin in [sys.__stdin__]:  # 简单实现，实际可优化
                try:
                    choice = input().strip()
                    if choice in songs:
                        change_song(choice)
                except:
                    pass
            time.sleep(0.08)

    except Exception as e:
        print(f"❌ 错误: {e}")
    finally:
        if serial:
            try:
                serial.close()
            except:
                pass

def graceful_shutdown():
    print("\n正在退出程序...")
    running.clear()
    time.sleep(0.6)
    pygame.mixer.quit()
    print("🎹 程序已安全退出")
    sys.stdout.flush()
    import os
    os._exit(0)

if __name__ == "__main__":
    thread = threading.Thread(target=start_brainlink, daemon=True)
    thread.start()

    print("程序运行中... 按 Ctrl + C 退出")

    try:
        while running.is_set():
            time.sleep(0.3)
    except KeyboardInterrupt:
        graceful_shutdown()
    except Exception:
        graceful_shutdown()